# -*- coding: UTF-8 -*-
# @File : server.py
# @IDE     : VScode
# @Author  : zhonggc
# @Date    : 2024/05/30 09:28:15

from fastapi import FastAPI
from langserve import add_routes
from format_chain import get_answer, get_final_chain, chat_history
from chat_history import ChatHistory
from starlette.responses import StreamingResponse

chat_history_hander = ChatHistory()

app = FastAPI(
    title="AI知识助手",
    description="AI知识助手",
    version="1.0",
)


# 初始化对话模型
chain = get_final_chain()
add_routes(
    app,
    chain,
    path="/chat",
)


@app.post("/chat")
async def chat(data: dict):
    input_text = data.get("input")
    session_id = data.get("session_id")
    # answer = get_answer(input_text=input_text, session_id=session_id)
    # 保存聊天记录到数据库
    # chat_history_hander.save_chat_history(session_id, chat_history.messages)
    return StreamingResponse(
        get_answer(input_text=input_text, session_id=session_id),
        media_type="text/event-stream",
    )


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="localhost", port=8000)
